package cas.ihep.evt.impl;

import cas.ihep.evt.EvtPDL;
import cas.ihep.pwa.evtgen.*;
import com.mongodb.client.*;
import com.mongodb.client.model.Filters;
import org.bson.Document;
import org.luaj.vm2.LuaTable;
import org.luaj.vm2.LuaValue;
import org.luaj.vm2.lib.OneArgFunction;

import javax.script.*;
import java.io.Closeable;
import java.io.FileReader;
import java.io.IOException;
import java.util.*;

public class MongoEvtPDL implements EvtPDL,Closeable{

    private MongoClient client;
    private MongoCollection<Document> collection;

    private transient HashMap<String,Document> cache;

    public MongoEvtPDL(String conn){
        client= MongoClients.create(conn);
        MongoDatabase database = client.getDatabase("evtpdl");
        collection= database.getCollection("EvtPDL");
        cache=new HashMap<>();
    }

    private Document findDocument(String evt){
//        Document doc=cache.get(evt);
//        if(doc==null) {
//            FindIterable<Document> docs = collection.find(Filters.eq(evt));
//            doc = docs.first();
//            assert doc != null;
//            cache.put(evt,doc);
//        }
//        return doc;
        FindIterable<Document> docs = collection.find(Filters.eq(evt));
        Document doc=docs.first();
        if(doc==null){
            throw new RuntimeException("No such event name "+evt);
        }
        return doc;
    }

    private Document findAlias(Document doc){
        String alias=doc.getString("alias");
        if(alias!=null && !alias.isEmpty()) {
            return findDocumentWithAlias(alias);
        }
        return doc;
    }

    private Document findDocumentWithAlias(String evt){
        FindIterable<Document> docs = collection.find(Filters.eq(evt));
        Document doc=docs.first();
        if(doc==null){
            throw new RuntimeException("No such event name "+evt);
        }
        String alias=doc.getString("alias");
        if(alias!=null && !alias.isEmpty()){
            return findDocument(alias);
        }
        return doc;
    }

    private double getDouble(Document doc,String field){
        return doc.getDouble(field);
    }

    private int getInt(Document doc,String field){
        return doc.getInteger(field);
    }

    private String getString(Document doc, String field){
        return doc.getString(field);
    }

    @Override
    public double getMeanMass(String evt) {
        return getDouble(findDocumentWithAlias(evt),"Mass");
    }

    @Override
    public double getMass(String evt) {
        // rollMass()
        Document doc= findDocumentWithAlias(evt);
        double width=doc.getDouble("Width"),mass=doc.getDouble("Mass");
        double massMin=doc.getDouble("MinMass"),massMax=doc.getDouble("MaxMass");
        double ymin, ymax;
        double temp;
        if ( width < 0.0001 ) {
            return mass;
        }else{
            ymin = Math.atan( 2.0*(massMin-mass)/width);
            ymax = Math.atan( 2.0*(massMax-mass)/width);
            temp= ( mass + ((width/2.0)*Math.tan(EvtRandom.flat(ymin,ymax))));
            return temp;
        }
    }

    @Override
    public double getMaxMass(String evt) {
        return getDouble(findDocumentWithAlias(evt),"MaxMass");
    }

    @Override
    public double getMinMass(String evt) {
        return getDouble(findDocumentWithAlias(evt),"MinMass");
    }

    @Override
    public double getMaxRange(String evt) {
        return getDouble(findDocumentWithAlias(evt),"MaxRange");
    }

    @Override
    public double getCtau(String evt) {
        return getDouble(findDocumentWithAlias(evt),"Ctau");
    }

    @Override
    public int getStdHep(String evt) {
        return getInt(findDocumentWithAlias(evt),"Stdhep");
    }

    @Override
    public int getLundKC(String evt) {
        return getInt(findDocumentWithAlias(evt),"LundKC");
    }

    @Override
    public int getChg3(String evt) {
        return getInt(findDocumentWithAlias(evt),"Chg3");
    }

    @Override
    public double get2ndMass(String evt) {
        return getDouble(findDocumentWithAlias(evt),"Mass2nd");
    }

    @Override
    public double get2ndWidth(String evt) {
        return getDouble(findDocumentWithAlias(evt),"Width2nd");
    }

    @Override
    public String getLStype(String evt) {
        return getString(findDocumentWithAlias(evt),"Lstype");
    }

    public String getChgconj(String evt){
        return getString(findDocument(evt),"Chgconj");
    }

    @Override
    public void reSetMass(String evt, double val) {
        Document doc=cache.get(evt);
        if(doc==null){
            doc=new Document("$set",new Document("Mass",val));
        }else{
            doc.put("Mass",val);
            doc=new Document("$set",doc);
        }
        collection.updateOne(Filters.eq(evt),doc);
    }

    @Override
    public void setLStype(String evt, String lstype) {
        Document doc=cache.get(evt);
        if(doc==null){
            doc=new Document("$set",new Document("Lstype",lstype));
        }else{
            doc.put("Lstype",lstype);
            doc=new Document("$set",doc);
        }
        collection.updateOne(Filters.eq(evt),doc);
    }

    @Override
    public void reSetWidth(String evt, double val) {
        Document doc= findDocumentWithAlias(evt);
        double _massMax,_massMin;
        double maxdelta=15.0*val;
        double _maxRange=doc.getDouble("MaxRange");
        double _mass=doc.getDouble("Mass");
        if ( _maxRange > 0.00001 ) {
            _massMin=_mass-_maxRange;
        }else{
            _massMin=_mass-15.0*val;
        }
        if ( _massMin< 0. ) _massMin=0.;
        _massMax=_mass+maxdelta;
        doc.put("Width", val);
        doc.put("MaxMass",_massMax);
        doc.put("MinMass",_massMin);
        doc=new Document("$set",doc);
        collection.updateOne(Filters.eq(evt),doc);
    }

    @Override
    public void reSetMaxMass(String evt, double val) {
        Document doc=cache.get(evt);
        if(doc==null){
            doc=new Document("$set",new Document("MaxMass",val));
        }else{
            doc.put("MaxMass",val);
            doc=new Document("$set",doc);
        }
        collection.updateOne(Filters.eq(evt),doc);
    }

    @Override
    public void reSetMinMass(String evt, double val) {
        Document doc=cache.get(evt);
        if(doc==null){
            doc=new Document("$set",new Document("MinMass",val));
        }else{
            doc.put("MinMass",val);
            doc=new Document("$set",doc);
        }
        collection.updateOne(Filters.eq(evt),doc);
    }

    @Override
    public void setVpars(String evt, List<Double> val) {
        Document doc=cache.get(evt);
        if(doc==null){
            doc=new Document("$set",new Document("Vpars",val));
        }else{
            doc.put("Vpars",val);
            doc=new Document("$set",doc);
        }
        collection.updateOne(Filters.eq(evt),doc);
    }

    @Override
    public void setVres(String evt, List<Double> val) {
        Document doc=cache.get(evt);
        if(doc==null){
            doc=new Document("$set",new Document("Vres",val));
        }else{
            doc.put("Vres",val);
            doc=new Document("$set",doc);
        }
        collection.updateOne(Filters.eq(evt),doc);
    }

    @Override
    public void setVstrs(String evt, List<String> val) {
        Document doc=cache.get(evt);
        if(doc==null){
            doc=new Document("$set",new Document("Vstrs",val));
        }else{
            doc.put("Vstrs",val);
            doc=new Document("$set",doc);
        }
        collection.updateOne(Filters.eq(evt),doc);
    }

    @Override
    public List<Double> getVpars(String evt) {
        Document doc= findDocument(evt);
        return (List<Double>)doc.get("Vpars");
    }

    @Override
    public List<Double> getVRes(String evt) {
        return (List<Double>) findDocument(evt).get("Vres");
    }

    @Override
    public List<String> getVstrs(String evt) {
        return (List<String>) findDocument(evt).get("Vstrs");
    }

    public double getWidth(String evt){
        return getDouble(findDocument(evt),"Width");
    }

    public String getSpintype(String evt){
        return getString(findDocument(evt),"Spintype");
    }

    @Override
    public void close() throws IOException {
        client.close();
    }

    private double checkGetDouble(LuaValue val){
        if (val.isnumber()){
            return val.todouble();
        }
        throw new RuntimeException(val.toString()+" is not a double value");
    }

    private String checkGetString(LuaValue val){
        if (val.isstring()){
            return val.tojstring();
        }
        throw new RuntimeException(val.toString()+" is not a String value");
    }

    private int checkGetInt(LuaValue val){
        if(val.isnumber()){
            return val.toint();
        }
        throw new RuntimeException(val.toString()+" is not a numeric value");
    }

    private void readPDT(ScriptEngine e, String filename)throws IOException,ScriptException {
        CompiledScript cs=((Compilable)e).compile(new FileReader(filename));
        Bindings bd=new SimpleBindings();
        bd.put("pdtadd", new OneArgFunction() {
            @Override
            public LuaValue call(LuaValue luaValue) {
                return pdtadd(luaValue);
            }
        });
        bd.put("pdtsets", new OneArgFunction() {
            @Override
            public LuaValue call(LuaValue luaValue) {
                return LuaValue.NIL;
            }
        });
        cs.eval(bd);
    }

    private LuaValue pdtadd(LuaValue luaValue){
        if(luaValue.istable()){
            LuaTable table=(LuaTable)luaValue;
            Document doc=new Document();
            String pname=checkGetString(table.get(3));
            doc.append("_id",pname);
            int stdhepid=checkGetInt(table.get(4));
            double mass=checkGetDouble(table.get(5)),pwidth=checkGetDouble(table.get(6)),pmaxwidth=checkGetDouble(table.get(7));
            int chg3=checkGetInt(table.get(8)),spin2=checkGetInt(table.get(9));
            double ctau=checkGetDouble(table.get(10));
            int lundkc=checkGetInt(table.get(11));
            if("string".equals(pname)){
                doc.append("Spintype","STRING");
            }else if("vpho".equals(pname)){
                doc.append("Spintype","VECTOR");
            }else {
                switch (spin2) {
                    case 0:
                        doc.append("Spintype", "SCALAR");
                        break;
                    case 1:
                        doc.append("Spintype", mass < 0.0001 ? "NEUTRINO" : "DIRAC");
                        break;
                    case 2:
                        doc.append("Spintype", mass < 0.0001 ? "PHOTON" : "VECTOR");
                        break;
                    case 3:
                        doc.append("Spintype", "RARITASCHWINGER");
                        break;
                    case 4:
                        doc.append("Spintype", "TENSOR");
                        break;
                    case 5:
                        doc.append("Spintype", "SPIN5HALF");
                        break;
                    case 6:
                        doc.append("Spintype", "SPIN3");
                        break;
                    case 7:
                        doc.append("Spintype", "SPIN7HALF");
                        break;
                    case 8:
                        doc.append("Spintype", "SPIN4");
                        break;
                    default:
                        doc.append("Spintype", "SCALAR");
                }
            }
            doc.append("Stdhep",stdhepid).append("LundKC",lundkc).append("Ctau",ctau).append("Chg3",chg3);
            absLineShape(doc,mass,pwidth,pmaxwidth);
            collection.insertOne(doc);
            return LuaValue.NIL;
        }
        return LuaValue.error("not a table");
    }

    private void absLineShape(Document doc,double mass,double width,double maxRange){
        double maxdelta=15.0*width;
        double _massMax,_massMin;
        //if ( width>0.001 ) {
        //  if ( 5.0*width < 0.6 ) maxdelta = 0.6;
        //}
        if ( maxRange > 0.00001 ) {
            _massMin=mass-maxRange;
        }
        else{
            _massMin=mass-15.0*width;
        }
        if ( _massMin< 0. ) _massMin=0.;
        _massMax=mass+maxdelta;
        doc.append("Mass", mass).append("Width", width).append("MaxRange", maxRange).append("MaxMass",_massMax).append("MinMass",_massMin);
    }

    void alias(String name,String newname){
        FindIterable<Document> id=collection.find(Filters.eq(name)),aliasid=collection.find(Filters.eq(newname));
        Iterator<Document> iditer=id.iterator();
        Iterator<Document> aliasiter=aliasid.iterator();
        if(iditer.hasNext() && !aliasiter.hasNext()){
            Document doc=new Document();
            doc.append("_id",newname).append("alias",name);
            collection.insertOne(doc);
        }
    }

    private String chargeConjNoAlias(Document doc,String evtname){
        if(evtname==null){
            evtname=doc.getString("_id");
        }
        int stdhep=-doc.getInteger("Stdhep");
        FindIterable<Document> results=collection.find(Filters.eq("Stdhep",stdhep));
        Document result=results.first();
        if(result!=null){
            String chargeConjName=result.getString("_id");
            collection.updateOne(Filters.eq(evtname),new Document("$set",new Document("Chgconj",chargeConjName)));
            return chargeConjName;
        }
        collection.updateOne(Filters.eq(evtname),new Document("$set",new Document("Chgconj",evtname)));
        return evtname;
    }

    private String chargeConj(String evtname){
        Document doc= findDocument(evtname);
        String alias=doc.getString("alias");
        if(alias==null || alias.isEmpty()){
            return chargeConjNoAlias(doc,evtname);
        }else if(evtname.equals(chargeConj(alias))){
            collection.updateOne(Filters.eq(evtname),new Document("$set",new Document("Chgconj",evtname)));
            return evtname;
        }else{
            throw new RuntimeException("Trying to charge conjugate alias particle: "+evtname+" without defining the alias!");
        }
    }

    void aliasChgConj(String a,String abar){
        Document doc=findDocument(a);
        Document origin=findAlias(doc);
        if(abar.equals(chargeConjNoAlias(origin,null))){
            throw new RuntimeException("Can't charge conjugate the two aliases: "+a+" and "+abar);
        }
        collection.updateOne(Filters.eq(a),new Document("$set",new Document("Chgconj",abar)));
        collection.updateOne(Filters.eq(abar),new Document("$set",new Document("Chgconj",a)));
    }

    private LuaValue Particle(LuaValue name,LuaValue newMass,LuaValue newWidth){
        if(name.isstring() && newMass.isnumber() && (newWidth.isnumber() || newWidth.isnil())){
            double newMassVal=newMass.todouble();
            String jname=name.tojstring();
            Document doc=findDocumentWithAlias(jname);
            double meanmass=doc.getDouble("Mass");
            double newWidthVal=newWidth.isnumber()?newWidth.todouble():meanmass;
            doc=new Document();
            doc.append("Width",newWidthVal).append("Mass",newMassVal);
            collection.updateOne(Filters.eq(jname),new Document("$set",doc));
            return LuaValue.NIL;
        }
        return LuaValue.error("Invalid type of parameters");
    }

    private LuaValue Vparameters(LuaValue name,LuaValue vpars){
        if(name.isstring() && vpars.istable()){
            LuaTable tvpars=(LuaTable)vpars;
            int sz=tvpars.keyCount();
            ArrayList<Double> vparslist=new ArrayList<>();
            for(int i=0;i<sz;i++){
                LuaValue ele=tvpars.get(i+1);
                if (!ele.isnumber()){
                    return LuaValue.error("Invalid type of parameters");
                }
                vparslist.add(ele.todouble());
            }
            collection.findOneAndUpdate(Filters.eq(name.tojstring()),new Document("$set",new Document("Vpars",vparslist)));
            return LuaValue.NIL;
        }
        return LuaValue.error("Invalid type of parameters");
    }

    private LuaValue Vstrings(LuaValue name,LuaValue table){
        if(name.isstring() && table.istable()){
            LuaTable vstrs=(LuaTable)table;
            int sz=vstrs.keyCount();
            ArrayList<String> vstrslist=new ArrayList<>();
            for(int i=0;i<sz;i++){
                LuaValue ele=vstrs.get(i+1);
                if(!ele.isstring()){
                    return LuaValue.error("Invalid type of parameter");
                }
                vstrslist.add(ele.tojstring());
            }
            collection.findOneAndUpdate(Filters.eq(name.tojstring()),new Document("$set",new Document("Vstrs",vstrslist)));
            return LuaValue.NIL;
        }
        return LuaValue.error("Invalid type of parameters");
    }

    private LuaValue changeLStype(LuaValue name,String type){
        if (name.isstring()){
            collection.findOneAndUpdate(Filters.eq(name.tojstring()),new Document("$set",new Document("Lstype",type)));
            return LuaValue.NIL;
        }
        return LuaValue.error("Invalid type of parameters");
    }

    private LuaValue Define(LuaValue name, LuaValue value, Map<String,String> symtable){
        if (name.isstring() && (value.isnumber() && value.isstring())){
            symtable.put(name.tojstring(),value.tojstring());
            return LuaValue.NIL;
        }
        return LuaValue.error("Invalid type of parameters");
    }

    private void reserveCapacity(List<EvtParticleDecayList> decaytable, int sz){
        int ssz=decaytable.size();
        for(;ssz<sz;ssz++){
            decaytable.add(new EvtParticleDecayList());
        }
    }

    private LuaValue Alias(LuaValue newname,LuaValue oldname,List<EvtParticleDecayList> decaytable){
        if(newname.isstring() && oldname.isstring()){
            alias(newname.tojstring(),newname.tojstring());
            reserveCapacity(decaytable,(int)collection.count());
            return LuaValue.NIL;
        }
        return LuaValue.error("Invalid type of parameters");
    }

    private LuaValue ChargeConj(LuaValue aname,LuaValue abarname){
        if(aname.isstring() && abarname.isstring()){
            aliasChgConj(aname.tojstring(),abarname.tojstring());
            return LuaValue.NIL;
        }
        return LuaValue.error("Invalid type of parameters");
    }

    private LuaValue ModelAlias(LuaValue aliasname,LuaValue modelname,LuaValue table,List<EvtModelAlias> modelAliasList){
        if(aliasname.isstring() && modelname.isstring() && table.istable()){
            ArrayList<String> modelarglist=new ArrayList<>();
            LuaTable malist=(LuaTable)table;
            int sz=malist.keyCount();
            for(int i=0;i<sz;i++){
                LuaValue val=malist.get(i+1);
                if(!val.isnumber() && !val.isstring()){
                    return LuaValue.error("Invalid type of parameters");
                }
                modelarglist.add(val.tojstring());
            }
            modelAliasList.add(new EvtModelAlias(aliasname.tojstring(),modelname.tojstring(),modelarglist));
            return LuaValue.NIL;
        }
        return LuaValue.error("Invalid type of parameters");
    }

    private LuaValue Decay(LuaValue parent, LuaValue table,
                           EvtModel modellist, List<EvtModelAlias> modelAliasList, List<EvtParticleDecayList> decaytable,
                           Map<String,String> symtable){
        if (parent.isstring() && table.istable()){
            reserveCapacity(decaytable,partlist.size());
            double brfrsum=0;
            EvtId ipar=getEvtId(parent.tojstring());
            if (decaytable.get(ipar.getAlias())!=null && decaytable.get(ipar.getAlias()).getNMode()!=0){
                decaytable.get(ipar.getAlias()).removeDecay();
            }
            LuaTable subvalues=(LuaTable)table;
            int sz=subvalues.keyCount();
            for(int i=0;i<sz;i++){
                LuaValue sv=subvalues.get(i+1);
                if(!sv.istable()){
                    return LuaValue.error("Invalid type of parameters");
                }
                LuaTable stable=(LuaTable)sv;
                int nextValue=1;
                LuaValue brfr=stable.get(nextValue++);
                if (!brfr.isnumber()){
                    return LuaValue.error("Invalid type of parameters");
                }
                LuaValue daugs=stable.get(nextValue++);
                LuaTable daugstable=(LuaTable)daugs;
                ArrayList<EvtId> daugpart=new ArrayList<>();
                int daugsz=daugstable.keyCount();
                String firstDaug=daugstable.get(1).tojstring();
                for(int j=0;j<daugsz;j++){
                    daugpart.add(getEvtId(daugstable.get(j+1).tojstring()));
                }
//                boolean isname=daugsz>0 && daugpart.get(0).getId()>=0;
//                int ismodel=modellist.isModel(firstDaug)?1:0;
//                if(!(isname || ismodel!=0)){
//                    int iAlias,__sz=modelAliasList.size();
//                    for(iAlias=0;iAlias<__sz;iAlias++){
//                        if ( modelAliasList.get(iAlias).matchAlias(firstDaug) ) {
//                            ismodel=2;
//                            break;
//                        }
//                    }
//                }
//                if (!(isname || ismodel!=0)){
//                    throw new RuntimeException();
//                }
                LuaValue luamodel=stable.get(nextValue++);
                String model=luamodel.tojstring();
                if("PHOTOS".equals(model)){
                    //System.out.println("PHOTOS");
                    model=stable.get(nextValue++).tojstring();
                }else if("VERBOSE".equals(model)){
                    //System.out.println("VERBOSE");
                    model=stable.get(nextValue++).tojstring();
                }else if("SUMMARY".equals(model)){
                    //System.out.println("SUMMARY");
                    model=stable.get(nextValue++).tojstring();
                }
                int iAlias;
                int foundAnAlias=-1;
                for(iAlias=0;iAlias<modelAliasList.size();iAlias++){
                    if ( modelAliasList.get(iAlias).matchAlias(model) ) {
                        foundAnAlias=iAlias;
                        break;
                    }
                }
                if(foundAnAlias==-1){
                    if(!modellist.isModel(model)){
                        throw new RuntimeException(model + " is not a valid model");
                    }
                }else{
                    model=modelAliasList.get(foundAnAlias).getName();
                }
                LuaTable argtable=(LuaTable)stable.get(nextValue);
                List<String> temp_fcn_new_args;
                if (foundAnAlias==-1){
                    temp_fcn_new_args=new ArrayList<>();
                    int keysize=argtable.keyCount();
                    for(int idx=0;idx<keysize;idx++){
                        temp_fcn_new_args.add(symGet(symtable,argtable.get(idx+1).tojstring()));
                    }
                }else{
                    temp_fcn_new_args=modelAliasList.get(foundAnAlias).getArgList();
                }
                double jbrfr=brfr.todouble();
                brfrsum+=jbrfr;
                EvtDecayBase temp_fcn_new=modellist.getFcn(model);
                double massmin=0;
                temp_fcn_new.saveDecayInfo(ipar,daugsz,daugpart,temp_fcn_new_args.size(),temp_fcn_new_args,model,jbrfr);
                for(EvtId did:daugpart){
                    if(getMinMass(did)>0.0001){
                        massmin+=getMinMass(did);
                    }else{
                        massmin+=getMeanMass(did);
                    }
                }
                decaytable.get(ipar.getAlias()).addMode(temp_fcn_new,brfrsum,massmin);
            }
            decaytable.get(ipar.getAlias()).finish();

            for(int iii=0;iii<partlist.size();iii++){
                int nModTot=decaytable.get(iii).getNMode();
                if(nModTot==0){
                    continue;
                }
                if(getWidth(iii)<0.0000001){
                    continue;
                }
                double minMass=getMaxMass(iii);
                for(int jjj=0;jjj<nModTot;jjj++){
                    double tmass=decaytable.get(iii).getDecay(jjj).getMassMin();
                    if(tmass<minMass){
                        minMass=tmass;
                    }
                }
                if(minMass>getMinMass(iii)){
                    reSetMinMass(iii,minMass);
                }
            }
            return LuaValue.NIL;
        }
        return LuaValue.error("Invalid type of parameters");
    }
}
